import torch
from load_data import train_dataloader, test_dataloader
import cv2
from setting import device

_net = torch.load('../model/net_epoch10.pkl').to(device)

_net.eval()  # 进入测试模式

acc = 0
for i, (x, y) in enumerate(test_dataloader):
    x = x.to(device)
    y = y.to(device)
    out = _net(x)

    pred = out.argmax(dim=1)
    acc += (pred == y).sum().item()
    for j in range(len(pred)):
        if pred[j] != y[j]:
            _x = (x[j].permute(1, 2, 0) * 255).int().to('cpu').numpy()
            cv2.imwrite(f'../img/bad_{j}.jpg', _x)
            print(j)
print('acc', acc)
